import torch
from torch.utils.data import Dataset

device = torch.device("cpu")


class StockDataset(Dataset):

    def __init__(self, dataFrame, win_size):
        self.dataFrame = dataFrame
        self.data = []
        self.labels = []
        self.win_size = win_size

        for i in range(self.win_size, len(self.dataFrame) - self.win_size):
            x = torch.tensor(self.dataFrame.iloc[i:i + self.win_size, :-1].values, dtype=torch.float).to(device)
            y = torch.tensor(self.dataFrame.iloc[i:i + self.win_size, -1:].values, dtype=torch.float).to(device)
            self.data.append(x)
            self.labels.append(y)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]
